package com.gitee.wsl.coroutines.channels.receive

import kotlinx.coroutines.channels.ClosedReceiveChannelException
import kotlinx.coroutines.channels.ReceiveChannel
import kotlinx.coroutines.channels.onClosed
import kotlinx.coroutines.channels.onFailure
import kotlinx.coroutines.channels.onSuccess
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kotlin.coroutines.cancellation.CancellationException

fun <T> ReceiveChannel<T>.toBatchedFlow(maxBatchSize: Int? = null): Flow<List<T>> = flow {
    while (true) {
        val batch = mutableListOf<T>()
        try {
            batch.add(this@toBatchedFlow.receive())
        } catch (e: ClosedReceiveChannelException) {
            return@flow
        } catch (e: CancellationException) {
            return@flow
        }
        for (i in 1..<(maxBatchSize ?: Int.MAX_VALUE)) {
            var isChannelCurrentlyEmpty = false
            this@toBatchedFlow.tryReceive()
                .onSuccess { batch.add(it) }
                .onFailure { isChannelCurrentlyEmpty = true }
                .onClosed {
                    this@flow.emit(batch)
                    return@flow
                }
            if (isChannelCurrentlyEmpty) break
        }
        this@flow.emit(batch)
    }
}
